require 'openssl'
require 'base64'

class Enigma

  class EncryptError < StandardError; end
  class DecryptError < StandardError; end

  # Creates an instance of Enigma using given shared secret
  # @raise ArgumentError if secret is not a string
  def initialize(secret)
    unless secret.is_a? String
      fail ArgumentError, 'secret must be a string'
    end
    @key = Digest::SHA1.hexdigest(secret)
  end

  # Encrypts a string.
  #
  # @param message a string to encrypt
  # @return a String
  # @raise ArgumentError if message is not a string
  # @raise EncryptError if encryption failed
  def encrypt(message)
    unless message.is_a? String
      fail ArgumentError, 'message must be a string'
    end

    begin
      cipher = OpenSSL::Cipher.new('aes-256-cbc')

      cipher.encrypt
      cipher.key = @key

      cipher.iv = iv = cipher.random_iv

      encrypted = cipher.update(message)
      encrypted << cipher.final

      pack(iv, encrypted)
    rescue => _
      fail EncryptError, 'Could not encode message.'
    end
  end

  # Decrypts an encrypted messages.
  # If the secret does not match the encrypted message then
  # an Enigma::DecryptError is raised.
  #
  # @param encrypted_message a string to decrypt (outputted by Enigma#encrypt)
  # @return a String with the decrypted message
  # @raise ArgumentError if encrypted_message is not a string
  # @raise DecryptError if Decryption failed

  def decrypt(encrypted_message)
    unless encrypted_message.is_a? String
      fail ArgumentError, 'encrypted_message must be a string'
    end

    begin
      iv, encrypted_message = unpack(encrypted_message)

      cipher = OpenSSL::Cipher.new('aes-256-cbc')
      cipher.decrypt
      cipher.key = @key

      cipher.iv = iv

      decrypted = cipher.update(encrypted_message)
      decrypted << cipher.final

      decrypted
    rescue => _
      fail DecryptError, 'Could not decode encrypted message.'
    end
  end

  protected

  # Packing the encrypted message into a string.
  #
  # Enigma#pack and Enigma#unpack can be overridden from a subclass.
  # Just make sure that the following still holds:
  #
  #     iv2, encrypted2 = unpack(pack(iv1, encrypted1)
  #     iv1 == iv2 and encrypted1 == encrypted2
  #
  # @param iv the random iv (bytes) generated by the cipher
  # @param encrypted the encrypted bytes
  # @return String with serialized iv and encrypted parameters
  def pack(iv, encrypted)
    [iv, encrypted].map { |x| Base64.encode64(x).strip }.join('|')
  end

  # Unpacking the encrypted_message into an array of iv and encrypted data.
  #
  # Enigma#pack and Enigma#unpack can be overridden from a subclass.
  # Just make sure that the following still holds:
  #
  #     iv2, encrypted2 = unpack(pack(iv1, encrypted1)
  #     iv1 == iv2 and encrypted1 == encrypted2
  #
  # @param encrypted_message as outputted by Enigma#pack
  # @return Array containing iv bytes and encrypted bytes
  def unpack(encrypted_message)
    encrypted_message.split('|').map { |m| Base64.decode64(m) }
  end

end

